/*
 * @Author: chyang0822 270917365@qq.com
 * @Date: 2024-12-10 16:27:49
 * @LastEditors: chyang0822 270917365@qq.com
 * @LastEditTime: 2024-12-11 15:43:54
 * @FilePath: /tensorrt_guide/swin-transformer/semantic-segmentation/UpsamplePlugin.h
 * @Description:
 *
 * Copyright (c) 2024 by ${git_name_email}, All Rights Reserved.
 */
#ifndef UPSAMPLE_PLUGIN_H
#define UPSAMPLE_PLUGIN_H

#include "NvInferPlugin.h"
#include <string>
#include <vector>

using namespace nvinfer1;

class UpsamplePlugin : public IPluginV2 {
  public:
    UpsamplePlugin(const std::string name, float scale_h, float scale_w);

    UpsamplePlugin(const std::string name, const void *data, size_t length);

    // It doesn't make sense to make UpsamplePlugin without arguments, so we delete default constructor.
    UpsamplePlugin() = delete;

    virtual int getNbOutputs() const noexcept override;

    Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) noexcept override;

    int initialize() noexcept override;

    void terminate() noexcept override;

    size_t getWorkspaceSize(int) const noexcept override { return 0; };

    virtual int enqueue(int batchSize, const void *const *inputs, void *const*outputs, void *workspace,
                cudaStream_t stream) noexcept override;

    size_t getSerializationSize() const noexcept override;

    void serialize(void *buffer) const noexcept override;

    void configureWithFormat(const Dims *inputDims, int nbInputs, const Dims *outputDims, int nbOutputs, DataType type,
                             PluginFormat format, int maxBatchSize)noexcept override;

    bool supportsFormat(DataType type, PluginFormat format) const noexcept override;

    const char *getPluginType() const noexcept override;

    const char *getPluginVersion() const noexcept override;

    void destroy() noexcept override;

    nvinfer1::IPluginV2 *clone() const noexcept override;

    void setPluginNamespace(const char *pluginNamespace) noexcept override;

    const char *getPluginNamespace() const noexcept override;

  private:
    const std::string mLayerName;
    bool mAlignCorners;
    float mScaleFactor_h;
    float mScaleFactor_w;
    size_t mInputVolume;
    Dims3 mInputShape;
    std::string mNamespace;
};

class UpsamplePluginCreator : public IPluginCreator {
  public:
    UpsamplePluginCreator();

    const char *getPluginName() const noexcept override;

    const char *getPluginVersion() const noexcept override;

    const PluginFieldCollection *getFieldNames()noexcept override;

    IPluginV2 *createPlugin(const char *name, const PluginFieldCollection *fc)noexcept override;

    IPluginV2 *deserializePlugin(const char *name, const void *serialData, size_t serialLength)noexcept override;

    void setPluginNamespace(const char *pluginNamespace)noexcept override;

    const char *getPluginNamespace() const noexcept override;

  private:
    static PluginFieldCollection mFC;
    static std::vector<PluginField> mPluginAttributes;
    std::string mNamespace;
};

#endif
